ck_tile: fix 2:4-sparse SWMMAC correctness on gfx1201/RDNA4 (3 bugs) + fail→pass repro - #3759

Thanks for digging this out and writing it up in this much detail — this had been sitting without a
review for three weeks, and the analysis in the description made it possible to check the claims
rather than guess at them. I have a gfx1201 box, so I ran it.

Short version: bug 1's fix holds up — I reproduced a failure on the base tree and, for what this
test measures, traced it to that one line. But the committed test can't be built from the PR alone, it passes on
the unfixed tree in the configuration that is buildable, and it never reaches the iu4 path that
bugs 2 and 3 are about. Separately, the changed code is not restricted to gfx1201, and I think that's
the thing to sort out before merge.

This is not a formal approval: I can only speak for gfx1201 and the int8 path.

Environment. AMD Radeon RX 9070 XT, gfx1201, amdgcn-amd-amdhsa--gfx1201; Linux 6.14.0-37,
amdgpu 6.19.14.31400100; container
rocm/pytorch:rocm7.14_ubuntu24.04_py3.12_pytorch_release_2.12.0, AMD clang 23.0.0git
(ROCm/llvm-project 46fcb339fb61). Trees compared: base
8fc1ac24e9bd7b431663a15e2122ce02c2979d37 — which is git merge-base of this PR's head and
develop, and base..head is exactly the three files in this PR — versus head
ac24ac28d662acf279478545cec541d4dde00f31. Identical test source and compile flags on both sides,
separate clean trees and build directories, three fresh processes per variant, identical stdout on
every repeat.

Two environment-side notes so the commands reproduce, neither of which is a problem with this PR:
that container needs --rocm-device-lib-path=<sdk>/lib/llvm/amdgcn/bitcode, and it ships only
libamdhip64.so.7, so a libamdhip64.so symlink is needed for the link step.

1. The repro can't be built from the PR alone

sparse_swmmac_correctness_repro.cpp:34 has

#include "real_a_tile.h"  // REAL_A_16x128: Quark int8 2:4 weight tile

outside the #ifdef USE_REAL_TILE at line 264, and real_a_tile.h isn't one of the PR's three files.
From a clean checkout the translation unit doesn't compile either way, so the max_abs_err=352
figures can't be checked by a reviewer.

I generated a substitute tile to get something running: same int8_t[16][128] shape, at most two
non-zeros per group of four along K, covering all six two-survivor position pairs, all four
single-survivor positions and the zero-survivor case, with values derived from (row, group, position) so that a misplaced survivor is more likely to show up numerically rather than cancel.
This is my input, not yours — nothing below should be compared against 352.

2. base fails, head passes

The harness wraps hipMalloc / hipMemcpy / hipMemset / hipDeviceSynchronize / hipFree in
HIP_CHECK_ERROR, and none tripped, so the failing numbers are a real kernel result rather than a
silent launch failure.

3. The whole effect is bug 1

Two more variants — base plus only the nonzero_elems true-zero default, and head with only that
line reverted:

The head − bug-1 stdout is identical to base's, byte for byte (diff on the run logs). So for the
failure this test measures, on these two commits, with this input and these three shapes, that one
line is both necessary and sufficient. I'm not claiming more than that — this doesn't establish
correctness over all valid inputs, types or architectures. The reasoning in your comment matches what
I see.

4. The test never reaches iu4, so bugs 2 and 3 get no numerical coverage from it

I disassembled the code objects of the kernels that actually launch, rather than grepping the
executable. The three SparseGemmKernel symbols (WaveTileK = 32 / 64 / 128) contain 1, 2 and 4
SWMMAC instructions respectively, zero v_wmma, zero v_mfma — and all seven are
v_swmmac_i32_16x16x32_iu8. There is no iu4 instruction anywhere in the binary.

That follows from the instantiation: the test uses SparseMmaPipeline<int8_t, int8_t, int32_t, ...>,
and int8_t picks up the generic numeric_traits::PackedSize == 1, so the if constexpr(PackedSize == 1) branch is taken, the packed-nibble path of bug 2 is compiled out, the SWAP + XOR-1 transform of
bug 3 is never reached, and TotalCompressedElems * MmaOp::APackedSize evaluates unchanged.

So the description reasons about all three bugs from the iu4 side, but the committed test only
exercises iu8. I haven't verified bugs 2 or 3 either — building an iu4 oracle independently of the
transform under test (nibble order, sign extension, logical vs packed K) is its own job, and reusing
your transform as the oracle would be circular.

What I could check is the shape side, with a probe that instantiates
SparseMmaPipeline<pk_int4_t, pk_int4_t, int32_t, ...> directly and launches it:

Both trees compile, both emit v_swmmac_i32_16x16x32_iu4 and v_swmmac_i32_16x16x64_iu4, and
hipDeviceSynchronize() returns success. That's consistent with your ISA reading — 8 index values
per lane for a K=32 iu4 tile — but it's a shape and codegen check, not a correctness one.

5. The buildable configuration passes on the unfixed tree — and the reason is the one you identified

Built without -DUSE_REAL_TILE, so the #else synthetic path runs. (The substitute header from
§1 is still needed even here, since the #include is unconditional — dropping the define alone does
not make the file compile.)

I expected that to be because the synthetic input has no under-filled groups, but that isn't it. I
replayed the exact host-side fill (mt19937(42), uniform_int_distribution(-8, 8), then
apply_sparse_pattern) and counted:

So under-filled groups are plentiful — the distribution includes 0 — and it still passes on the buggy
code. The reason is exactly the precondition you call out in the comment: apply_sparse_pattern
always zeroes slots 1 and 3, so survivors only ever sit at slots 0 and 2, and the old
{a_vec[i*4+2], a_vec[i*4+3]} default therefore always seeds slot 1 from a guaranteed zero. My tile
breaks it because survivors also sit at positions 2 and 3.

The practical consequence: once the missing header is supplied in whatever form, the configuration
that does not define USE_REAL_TILE reports a green run on the bug. Worth committing an input (or
generating one in the test) that puts survivors at position 3, and making that the default.

6. The change isn't scoped to gfx1201, and bug 1's fix demonstrably reaches CDNA

This is my main question before merge, and I don't think it's answerable from gfx1201 alone.

Neither changed hunk has an architecture predicate; both branch on PackedSize / APackedSize.
More to the point, the transforms selector in sparse_transforms.hpp:381 is specialized purely on
the op family:

struct MmaTransformsDefaultSelector<MmaOp, CompilerTarget,
                                    std::enable_if_t<MmaOp::OpFamily == MmaOpFamily::SPARSE>>
{ using SelectedTransforms = MmaDefaultTransformsSparse<MmaOp::kCompressionRatio>; };

with no enable_if_target_family_gfx*, unlike the dense gfx9 / gfx11 / gfx12 selectors right
alongside it. So every SPARSE op on every target resolves to this compress_a_impl. And
sparse/mfma/sparse_gfx9.hpp defines SPARSE ops for GFX942 / GFX950 over fp16_t, bf16_t,
int8_t, fp8_t, bf8_t — none of which specialize PackedSize, so they all take the same
PackedSize == 1 branch that bug 1's fix changes.

I checked that rather than inferring it. A compile-only probe that instantiates
SparseMmaPipeline<int8_t, int8_t, int32_t, 16, 16, 64, ..., Gfx942Target> and asserts

static_assert(std::is_same_v<
    typename MmaTransformsDefaultSelector<MmaOp942, Gfx942Target>::SelectedTransforms,
    MmaDefaultTransformsSparse<MmaOp942::kCompressionRatio>>);

compiles on both trees at --offload-arch=gfx942, and the resulting code object contains
v_smfmac_i32_16x16x64_i8. So that CDNA instantiation does route through the compress_a_impl this
PR changes, and takes the PackedSize == 1 branch.

That makes bug 1's fix a behaviour change for gfx942 as well, for inputs whose survivors don't sit
where the old default assumed. I think it's the correct fix and my gfx1201 result supports it — but
whether CDNA results actually change in practice, and whether anything downstream depended on the old
behaviour, is a runtime question I can't answer without the hardware.

The narrower half of this: I found no pk_int4_t SPARSE op for gfx9 in the tree, so bug 2's packed
branch and bug 3's SWAP + XOR-1 look gfx12-only in practice today, which limits the blast radius of
the empirically-derived metadata mapping considerably. It's still worth saying out loud that those
branches carry no architecture condition, so a future CDNA pk_int4 sparse op would silently inherit a
mapping that was measured on gfx1201.

On CI: the GitHub API reports no check suites and no commit statuses for either cd34d2f or
ac24ac2. I can't tell from outside whether anything ran elsewhere, but from the PR there's no
visible regression signal for the architectures the selector reaches. That's a project-side
infrastructure question rather than something to put on you.

7. Minor, non-blocking

• The comments carry internal process labels (Stage-17b fix (local, not upstreamed),
  Stage-17c CLOSURE fix, this Stage-17c audit's own probe) — worth rewording for upstream.
• The new static_assert(PackedSize == 2 && is_same_v<LogicalADataType, pk_int4_t>) narrows the
  packed path to pk_int4_t explicitly. I haven't tested whether any other packed A type reaches
  this code today, so I don't know if it changes anything in practice — just noting it's a scope
  change beyond the three bugs.
• You already note the gtest port is pending. For what it's worth, test/ck_tile/CMakeLists.txt
  enumerates its subdirectories with explicit add_subdirectory(...) calls rather than a glob, and
  gfx1201_sparse_swmmac isn't among them, so the port will need that line too.

What would let me say more

1. Commit real_a_tile.h, or generate an equivalent tile inside the test, so the fail→pass is
   reproducible from the PR alone — and make sure the default-built configuration is one that fails
   before the fix.
2. Add an iu4 / pk_int4_t case, since bugs 2 and 3 have no numerical coverage without one.
3. Say what architecture scope is intended, given §6.

Happy to re-run any of this on gfx1201. If it would help, I can also try to pin bug 3's mapping
empirically and independently of sparse_transforms — a minimal wave-level kernel issuing raw
v_swmmac_*_iu4 with one-hot operands, enumerating the metadata encodings and reading the
contributing positions back out of the result — and then compare what that gives against SWAP + XOR-1.
That would be a second measurement rather than a derivation from the doc, which as you note doesn't
capture the iu4-specific detail. Say the word and I'll put the generator, flags, logs and per-symbol
disassembly somewhere you can pull them.